# -*- coding: utf-8 -*-
"""
省份地理边界可视化模块

此模块提供用于可视化中国省份地理边界数据的函数和类。
主要功能包括：
- 创建省份边界地图
- 生成包含省份边界数据和地图的HTML报告
- 绘制省份边界、网格点和干旱指数地图
"""

import os
import numpy as np
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from shapely.geometry import Point
from typing import List, Dict, Any, Optional, Tuple, Union


def create_province_map(gdf: gpd.GeoDataFrame, filtered_gdf: gpd.GeoDataFrame, 
                       output_path: str, name_field: str = 'NAME', 
                       figsize: Tuple[int, int] = (15, 10), dpi: int = 300) -> str:
    """
    创建省份边界地图并保存为图像文件
    
    参数:
        gdf: 包含所有省份数据的GeoDataFrame
        filtered_gdf: 筛选后的目标省份GeoDataFrame
        output_path: 输出图像文件路径
        name_field: 省份名称字段，默认为'NAME'
        figsize: 图像大小，默认为(15, 10)
        dpi: 图像分辨率，默认为300
        
    返回:
        保存的图像文件路径
    """
    print("正在生成可视化地图...")
    
    # 设置图形大小
    plt.figure(figsize=figsize)
    
    # 绘制所有省份作为背景（灰色）
    gdf.plot(ax=plt.gca(), color='lightgrey', edgecolor='grey', linewidth=0.5)
    
    # 绘制目标省份（彩色）
    filtered_gdf.plot(ax=plt.gca(), column=name_field, cmap='tab20', 
                     edgecolor='black', linewidth=0.8, legend=True)
    
    # 添加省份标签
    for idx, row in filtered_gdf.iterrows():
        # 获取几何中心点
        centroid = row['geometry'].centroid
        plt.text(centroid.x, centroid.y, row[name_field], 
                 fontsize=8, ha='center', va='center')
    
    # 设置标题和轴标签
    plt.title('指定省份地理位置和边界', fontsize=16)
    plt.xlabel('经度', fontsize=12)
    plt.ylabel('纬度', fontsize=12)
    
    # 添加网格线
    plt.grid(True, linestyle='--', alpha=0.7)
    
    # 保存图像
    plt.savefig(output_path, dpi=dpi, bbox_inches='tight')
    print(f"地图已保存到: {output_path}")
    
    # 关闭图形，释放内存
    plt.close()
    
    return output_path


def generate_html_report(result_df: pd.DataFrame, map_path: str, output_path: str) -> str:
    """
    生成包含省份边界数据和地图的HTML报告
    
    参数:
        result_df: 包含省份边界数据的DataFrame
        map_path: 地图图像文件路径
        output_path: 输出HTML文件路径
        
    返回:
        保存的HTML文件路径
    """
    html_content = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <meta charset="UTF-8">
        <title>省份经纬度范围报告</title>
        <style>
            body {{ font-family: Arial, sans-serif; margin: 20px; }}
            h1 {{ color: #2c3e50; }}
            table {{ border-collapse: collapse; width: 100%; margin-top: 20px; }}
            th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
            th {{ background-color: #f2f2f2; }}
            tr:nth-child(even) {{ background-color: #f9f9f9; }}
            .map-container {{ margin-top: 30px; text-align: center; }}
            img {{ max-width: 100%; border: 1px solid #ddd; }}
        </style>
    </head>
    <body>
        <h1>省份经纬度范围报告</h1>
        
        <h2>经纬度边界数据</h2>
        <table>
            <tr>
                <th>省份</th>
                <th>全称</th>
                <th>最小经度</th>
                <th>最小纬度</th>
                <th>最大经度</th>
                <th>最大纬度</th>
                <th>经度范围</th>
                <th>纬度范围</th>
            </tr>
    """
    
    # 添加表格行
    for idx, row in result_df.iterrows():
        html_content += f"""
            <tr>
                <td>{row['省份']}</td>
                <td>{row['全称']}</td>
                <td>{row['最小经度']:.6f}</td>
                <td>{row['最小纬度']:.6f}</td>
                <td>{row['最大经度']:.6f}</td>
                <td>{row['最大纬度']:.6f}</td>
                <td>{row['经度范围']}</td>
                <td>{row['纬度范围']}</td>
            </tr>
        """
    
    # 添加地图和结束标签
    html_content += f"""
        </table>
        
        <div class="map-container">
            <h2>省份地理位置地图</h2>
            <img src="{os.path.basename(map_path)}" alt="省份地图">
        </div>
    </body>
    </html>
    """
    
    # 保存HTML文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(html_content)
    
    print(f"HTML报告已保存到: {output_path}")
    
    return output_path


def plot_province_boundary(province_gdf: gpd.GeoDataFrame, province: str, 
                          name_field: str = 'NAME', figsize: Tuple[int, int] = (10, 8),
                          title: str = None, ax: Axes = None) -> Tuple[Figure, Axes]:
    """
    绘制省份边界
    
    参数:
        province_gdf: 包含省份边界数据的GeoDataFrame
        province: 省份名称
        name_field: 省份名称字段，默认为'NAME'
        figsize: 图像大小，默认为(10, 8)
        title: 图像标题，默认为None
        ax: matplotlib轴对象，默认为None
        
    返回:
        matplotlib图形和轴对象
    """
    # 筛选指定省份
    filtered_gdf = province_gdf[province_gdf[name_field] == province]
    
    if len(filtered_gdf) == 0:
        raise ValueError(f"未找到省份: {province}")
    
    # 创建图形和轴对象
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = ax.figure
    
    # 绘制省份边界
    filtered_gdf.plot(ax=ax, color='lightblue', edgecolor='black', linewidth=1.0)
    
    # 设置标题
    if title is None:
        title = f"{province}省份边界"
    ax.set_title(title, fontsize=14)
    
    # 设置轴标签
    ax.set_xlabel('经度', fontsize=12)
    ax.set_ylabel('纬度', fontsize=12)
    
    # 添加网格线
    ax.grid(True, linestyle='--', alpha=0.7)
    
    return fig, ax


def plot_grid_points(province_gdf: gpd.GeoDataFrame, province: str, grid_points: np.ndarray,
                    name_field: str = 'NAME', figsize: Tuple[int, int] = (10, 8),
                    title: str = None, ax: Axes = None) -> Tuple[Figure, Axes]:
    """
    绘制省份边界和网格点
    
    参数:
        province_gdf: 包含省份边界数据的GeoDataFrame
        province: 省份名称
        grid_points: 网格点坐标数组，形状为(n, 2)，每行为[经度, 纬度]
        name_field: 省份名称字段，默认为'NAME'
        figsize: 图像大小，默认为(10, 8)
        title: 图像标题，默认为None
        ax: matplotlib轴对象，默认为None
        
    返回:
        matplotlib图形和轴对象
    """
    # 先绘制省份边界
    fig, ax = plot_province_boundary(province_gdf, province, name_field, figsize, title, ax)
    
    # 绘制网格点
    if len(grid_points) > 0:
        ax.scatter(grid_points[:, 0], grid_points[:, 1], c='red', s=5, alpha=0.7, label='网格点')
        ax.legend()
    
    return fig, ax


def plot_drought_index_map(province_gdf: gpd.GeoDataFrame, province: str, 
                          grid_points: np.ndarray, index_values: np.ndarray,
                          name_field: str = 'NAME', figsize: Tuple[int, int] = (12, 10),
                          title: str = None, cmap: str = 'RdYlBu_r',
                          vmin: float = None, vmax: float = None,
                          ax: Axes = None) -> Tuple[Figure, Axes]:
    """
    绘制干旱指数地图
    
    参数:
        province_gdf: 包含省份边界数据的GeoDataFrame
        province: 省份名称
        grid_points: 网格点坐标数组，形状为(n, 2)，每行为[经度, 纬度]
        index_values: 干旱指数值数组，长度为n
        name_field: 省份名称字段，默认为'NAME'
        figsize: 图像大小，默认为(12, 10)
        title: 图像标题，默认为None
        cmap: 颜色映射，默认为'RdYlBu_r'
        vmin: 颜色映射最小值，默认为None
        vmax: 颜色映射最大值，默认为None
        ax: matplotlib轴对象，默认为None
        
    返回:
        matplotlib图形和轴对象
    """
    # 先绘制省份边界
    fig, ax = plot_province_boundary(province_gdf, province, name_field, figsize, title, ax)
    
    # 绘制干旱指数散点图
    if len(grid_points) > 0 and len(index_values) == len(grid_points):
        scatter = ax.scatter(grid_points[:, 0], grid_points[:, 1], 
                           c=index_values, cmap=cmap, s=10, alpha=0.8,
                           vmin=vmin, vmax=vmax)
        
        # 添加颜色条
        cbar = fig.colorbar(scatter, ax=ax, pad=0.01)
        cbar.set_label('干旱指数值', fontsize=12)
    
    # 设置标题
    if title is None:
        title = f"{province}干旱指数分布图"
    ax.set_title(title, fontsize=14)
    
    return fig, ax


def process_province_data(shp_path: str, target_provinces: List[str], 
                         output_dir: str = None, 
                         csv_filename: str = 'province_bounds_detailed.csv',
                         map_filename: str = 'province_map.png',
                         html_filename: str = 'province_report.html') -> Dict[str, str]:
    """
    处理省份数据的主函数，包括读取shapefile、提取边界、创建地图和生成报告
    
    参数:
        shp_path: shapefile文件路径
        target_provinces: 目标省份名称列表
        output_dir: 输出目录，默认为None（使用当前目录）
        csv_filename: CSV文件名，默认为'province_bounds_detailed.csv'
        map_filename: 地图文件名，默认为'province_map.png'
        html_filename: HTML报告文件名，默认为'province_report.html'
        
    返回:
        包含输出文件路径的字典
    """
    from .province_utils import read_province_shapefile, filter_province, extract_province_bounds
    
    # 设置输出目录
    if output_dir is None:
        output_dir = os.path.dirname(shp_path)
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 设置输出文件路径
    csv_path = os.path.join(output_dir, csv_filename)
    map_path = os.path.join(output_dir, map_filename)
    html_path = os.path.join(output_dir, html_filename)
    
    try:
        # 读取shapefile
        gdf = read_province_shapefile(shp_path)
        
        # 筛选目标省份
        filtered_gdf = filter_province(gdf, target_provinces)
        
        # 提取边界数据
        result_df = extract_province_bounds(filtered_gdf)
        
        # 保存CSV文件
        result_df.to_csv(csv_path, index=False, encoding='utf-8-sig')
        print(f"省份边界数据已保存到: {csv_path}")
        
        # 打印结果
        print("\n省份经纬度范围:")
        for idx, row in result_df.iterrows():
            print(f"{row['省份']} ({row['全称']}): 经度 {row['经度范围']}, 纬度 {row['纬度范围']}")
        
        # 创建地图
        create_province_map(gdf, filtered_gdf, map_path)
        
        # 生成HTML报告
        generate_html_report(result_df, map_path, html_path)
        
        print("\n处理完成!")
        
        return {
            'csv_path': csv_path,
            'map_path': map_path,
            'html_path': html_path
        }
        
    except Exception as e:
        print(f"处理数据时出错: {e}")
        raise